#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/18
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
# 代码6-1
# 读取数据
# 数据加载
from random import randint
import numpy as np
import torch

torch.set_default_tensor_type(torch.FloatTensor)
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision.utils import save_image
import shutil
import cv2
import random
from PIL import Image
import itertools


def to_img(x):
    """
    将输入的张量x转换为图像
    参数:
    x: 输入的张量
    返回值:
    out: 转换后的图像
    """
    # 进行归一化处理，将x的值范围从[-1,1]转换为[0,1]
    out = 0.5 * (x + 1)
    # 限制out的值在[0,1]范围内
    out = out.clamp(0, 1)
    # 将out的形状从一维变为四维，形状为(-1, 3, 256, 256)
    out = out.view(-1, 3, 256, 256)
    return out


# 获取data6文件夹的绝对路径
data_path = os.path.abspath('data6')
# 设定图像大小为256
image_size = 256
# 设定批处理大小为1
batch_size = 1

# 定义图像转换的流程
# 使用torchvision.transforms库中的Compose函数，将一系列图像处理操作组合起来
# 首先，使用Resize函数将图像大小调整到原来的1.12倍，使用BICUBIC插值方法，可以得到更好的图像质量
# 然后，使用RandomCrop函数随机裁剪出一个指定大小的图像
# 接着，使用RandomHorizontalFlip函数随机水平翻转图像，增加数据多样性
# 然后，使用ToTensor函数将PIL Image或者 ndarray 转换为tensor，并且归一化至[0, 1]
# 最后，使用Normalize函数对图像进行标准化处理，减小输入到模型中的方差，提高模型的收敛速度和效果
transform = transforms.Compose([transforms.Resize(int(image_size * 1.12), Image.BICUBIC),
                                transforms.RandomCrop(image_size),
                                transforms.RandomHorizontalFlip(),
                                transforms.ToTensor(),
                                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])


def _get_train_data(batch_size=1):

    # 适用于win系统，mac系统使用"/"
    # train_a_filepath = data_path + '\\trainA\\'
    # train_b_filepath = data_path + '\\trainB\\'

    train_a_filepath = data_path + os.sep + 'trainA' + os.sep
    train_b_filepath = data_path + os.sep + 'trainB' + os.sep

    train_a_list = os.listdir(train_a_filepath)
    train_b_list = os.listdir(train_b_filepath)

    train_a_result = []
    train_b_result = []

    # 使用random.sample函数从train_a_list中随机抽取batch_size个不重复的元素，生成一个新的列表numlist
    numlist = random.sample(range(0, len(train_a_list)), batch_size)

    for i in numlist:
        a_filename = train_a_list[i]
        a_img = Image.open(train_a_filepath + a_filename).convert('RGB')
        res_a_img = transform(a_img)
        train_a_result.append(torch.unsqueeze(res_a_img, 0))

        b_filename = train_b_list[i]
        b_img = Image.open(train_b_filepath + b_filename).convert('RGB')
        res_b_img = transform(b_img)
        train_b_result.append(torch.unsqueeze(res_b_img, 0))

    return torch.cat(train_a_result, dim=0), torch.cat(train_b_result, dim=0)


if __name__ == "__main__":
    # 调用 _get_train_data 函数获取训练数据
    train_a, train_b = _get_train_data(batch_size=1)

    # 打印训练数据的形状
    print("Train A shape:", train_a.shape)
    print("Train B shape:", train_b.shape)

    # 检查数据是否正确加载
    assert train_a.shape == (1, 3, 256, 256), "Train A shape is incorrect"
    assert train_b.shape == (1, 3, 256, 256), "Train B shape is incorrect"

    # 将张量转换为图像并保存
    img_a = to_img(train_a)
    img_b = to_img(train_b)

    # 会保存两张图片
    save_image(img_a, 'tmp' + os.sep + '6-test_a.png')
    save_image(img_b, 'tmp' + os.sep + '6-test_b.png')

    print("Test passed! Images saved as 6-test_a.png  and 6-test_b.png")
